import json
import time
from concurrent.futures import ThreadPoolExecutor

from Agent.FuncSolver import FuncSolver
from Agent.Actuator import ORMAgent
from Agent.Controller import Controller
from get_sample import get_func_sample
from utils.rewriter import rewriter
from FantasticSQL import FantasticSQL
from FantasticORMnsExp import FantasticORM
from FantasticORMnsExp2 import FantasticORM2

func_sample = get_func_sample()


class Solver:
    def __init__(self):
        self.agent = ORMAgent()
        self.agent.start_code_kernel()

    def solve(self, question_info):
        return_answer = ''
        res = FuncSolver(question_info['question'], func_sample)
        if res:
            return_answer += res + "\n"
            # return return_answer
        res = self.agent.run(question_info)
        if res:
            return_answer += res
        return return_answer


def save2json(save_data, save_path):
    with open(save_path, 'w', encoding='utf8') as f:
        f.write(json.dumps(save_data, ensure_ascii=False, indent=4))


def process_data(question_group):
    try:
        bad_answer_list = [i['answer_for_summary'] for i in question_group['team'] if
                           i['answer_for_summary'] == '' or i['answer_for_summary'] == 'None']
        if len(bad_answer_list) >= 1:
            main_agent = Controller(Solver())
            for question_inx in range(len(question_group['team'])):
                question_info = question_group['team'][question_inx]
                if question_inx == 0:
                    answer, log = main_agent.run(question_info, True)
                    question_info['answer_for_summary'] += str(answer)  # + str(log)
                    save2json(data, answer_path)
                else:
                    answer, log = main_agent.run(question_info, False)
                    question_info['answer_for_summary'] += str(answer)  # + str(log)
                    save2json(data, answer_path)
    except Exception as e:
        print(e)


def main():
    # 创建线程池
    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = [executor.submit(process_data, sub_data) for sub_data in data]
        for future in futures:
            future.result()  # 等待所有线程完成


def process_data_SQL(question_group):
    try:
        # 找到没解答的题目
        FantasticSQL(question_group)

    except Exception as e:
        print(e)


def main_SQL():
    # 创建线程池
    with ThreadPoolExecutor(max_workers=9) as executor:
        futures = [executor.submit(process_data_SQL, sub_data) for sub_data in data]
        for future in futures:
            future.result()  # 等待所有线程完成


def process_data_ORM(question_group):
    try:
        # 找到没解答的题目
        FantasticORM(question_group)

    except Exception as e:
        print(e)


def main_ORM():
    # 创建线程池
    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = [executor.submit(process_data_ORM, sub_data) for sub_data in data]
        for future in futures:
            future.result()  # 等待所有线程完成


def process_data_ORM2(question_group):
    try:
        # 找到没解答的题目
        FantasticORM2(question_group)

    except Exception as e:
        print(e)


def main_ORM2():
    # 创建线程池
    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = [executor.submit(process_data_ORM2, sub_data) for sub_data in data]
        for future in futures:
            future.result()  # 等待所有线程完成


if __name__ == '__main__':
    print('开始运行')
    # time.sleep(1)
    # in_param_path = sys.argv[1]
    # print(in_param_path)

    question_path = r'E:\v3zl\app\devlop_home\submit\金融复赛a榜.json'
    print('question_path: ', question_path)
    time.sleep(1)
    out_path = 'answers.json'
    answer_path = out_path
    with open(question_path, 'r', encoding='utf8') as f:
        data = json.load(f)

    for i in data:
        for j in i['team']:
            j['answer'] = ''
            j['answer_for_summary'] = ''

    start = time.time()
    main_SQL()

    if time.time() - start < 60 * 60 * 10:
        main_ORM()
    if time.time() - start < 60 * 60 * 10:
        main_ORM2()
    if time.time() - start < 60 * 60 * 10:
        main()

    with open(answer_path, 'w', encoding='utf8') as f:
        f.write(json.dumps(data, ensure_ascii=False, indent=4))

    rewriter(answer_path)
